package com.base.queue;

import java.util.*;


public class MyPriorityQueue<E> {
	E[] arr;
	int cap;
	int size;
	private final Comparator<? super E> comparator;

	public boolean isEmpty() {
		return size == 0;
	}
	
	public E getTop() {
		return arr[0];
	}
	
	public int size() {
		return this.size;
	}
	

	public MyPriorityQueue(int cap, Comparator<? super E> comparator) {
		this.comparator = comparator;
		this.size = 0;
		this.cap = cap;
		this.arr = (E[]) new Object[cap];
	}
	
	private int getParentIndex(int x) {
		return (x - 1) / 2;
	}
	
	private void swap(int i, int j) {
		E t = arr[i];
		arr[i] = arr[j];
		arr[j] = t;
	}
	
	
	public void insert(E element) {
		if (size >= cap) return;
		arr[size] = element;
		int i = size;
		while(comparator.compare(arr[getParentIndex(i)], arr[i]) > 0) {
			int idx = getParentIndex(i);
			swap(idx, i);
			i = idx;
		}
		size ++;
	}
	
	public E poll() {
		E topEl = arr[0];
		arr[0] = arr[size - 1];
		adjust(0);
		size --;
		return topEl;
	}
	
	private void adjust(int i) {
		int left = i * 2 + 1;
		int right = i * 2 + 2;
		
		int idx = i;
		if (left < size && comparator.compare(arr[left], arr[idx]) < 0) {
			idx = left;
		}
		if (right < size && comparator.compare(arr[right], arr[idx]) < 0) {
			idx = right;
		}
		
		if (i != idx) {
			swap(idx, i);
			adjust(idx);
		}
	}
}

class Solution {
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> mp = new HashMap<>();
        for (int i = 0; i < nums.length; i ++) {
        	mp.put(nums[i], mp.getOrDefault(nums[i], 0) + 1);
        }
        
        MyPriorityQueue<Integer> queue = new MyPriorityQueue<>(k, new Comparator<Integer>() {

			@Override
			public int compare(Integer o1, Integer o2) {
				return mp.get(o1) - mp.get(o2);
			}
		});
        
        for (int key : mp.keySet()) {
        	if (queue.size() == k) {
        		if (mp.get(queue.getTop()) < mp.get(key)) {
        			queue.poll();
        			queue.insert(key);
        		}
        	} else {
        		queue.insert(key);
        	}
        }
        
        int[] res = new int[k];
        int i = 0;
        while(!queue.isEmpty()) {
        	res[i] = queue.poll();
        	i ++;
        }

        return res;
    }
}











